# -*- coding:utf-8 -*-
import backtrader as bt
#####################
import pandas as pd
import os
import datetime
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']

class MyStrategy(bt.Strategy):
    params = (
        ('maperiod_3', 3),
        ('maperiod_5',5)
    )

    def __init__(self):
        print(f'init___{self.datas[0].datetime.date(0)}')
        self.dataclose = self.datas[0].close
        self.sma_3 = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod_3)
        self.sma_5 = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod_5)
        self.order = None
        self.flag = 'sell'
        # self.sma_3.plotinfo.plot = False
        # self.sma_5.plotinfo.plot = False

    def start(self):
        pass

    def prenext(self):
        pass

    def nextstart(self):
        pass

    def next(self):
        if self.sma_3[0] > self.sma_5[0] and self.flag == 'sell':
            self.order = self.buy(size=1000)
            self.flag = 'buy'
            print(f"{self.datas[0].datetime.date(0)},买入！价格为{self.dataclose[0]}")
        else:
            if self.sma_3[0] < self.sma_5[0] and self.flag == 'buy':
                self.order = self.sell(size=1000)
                self.flag = 'sell'
                print(f"{self.datas[0].datetime.date(0)},卖出！价格为{self.dataclose[0]}")

    def stop(self):
        pass


if __name__ == '__main__':
    cerebro = bt.Cerebro()
    cerebro.broker.set_cash(100000.00)  # 设置初始资金金额
    cerebro = bt.Cerebro(stdstats=False)
    cerebro.addobserver(bt.observers.Broker)
    # cerebro.addobserver(bt.observers.Trades)
    cerebro.addobserver(bt.observers.BuySell)
    # cerebro.addobserver(bt.observers.DrawDown)
    cerebro.addobserver(bt.observers.Value)
    # cerebro.addobserver(bt.observers.TimeReturn)

    init_fund = cerebro.broker.getvalue()
    print(f'初始资金:{init_fund}')

    filename = 'sz000002.csv'
    print(filename)
    data = pd.read_csv(filename, index_col="date", parse_dates=True)
    # print(data)
    # 也可以如下处理日期
    # data.index=pd.to_datetime(data.date)
    # data.drop(columns=["date"],inplace=True)
    # print(data)

    load_data = bt.feeds.PandasData(dataname=data, fromdate=datetime.datetime(2020, 1, 1), todate=datetime.datetime(2020, 10, 18))
    cerebro.adddata(load_data)
    cerebro.addstrategy(MyStrategy)

    #####################
    cerebro.run()
    end_fund = cerebro.broker.getvalue()
    print(f'期末资金:{end_fund}')

    """
    处理错误
    cannot import name 'warnings' from 'matplotlib.dates
    pip uninstall matplotlib
    pip install matplotlib==3.2.2
    """
    cerebro.plot(style="candle")
